//
//  audio_encoder.cpp
//  Recoder
//
//  Created by dg on 2023/12/13.
//

#include "audio_encoder.hpp"

#define LOG_TAG "AudioEncoder"

extern "C" {
#include "libavutil/samplefmt.h"
}

bool AudioEncoder::checkSampleFMT(const AVCodec *codec, enum AVSampleFormat fmt) {
    
    const enum AVSampleFormat *p = codec->sample_fmts;
    while (*p != AV_SAMPLE_FMT_NONE) { // AV_SAMPLE_FMT_NONE 代表指向了最后一项
        if (*p == fmt) {
            return true;
        }
        p++;
    }
    
    return false;
}

int AudioEncoder::bestSampleRate(const AVCodec *codec) {
    const int *p;
    int best_samplerate = 0;

    if (!codec->supported_samplerates)
        return 44100;

    p = codec->supported_samplerates;
    while (*p) {
        if (!best_samplerate || abs(44100 - *p) < abs(44100 - best_samplerate))
            best_samplerate = *p;
        p++;
    }
    return best_samplerate;
}

int AudioEncoder::selectChannelLayout(const AVCodec *codec)
{
    const AVChannelLayout *p;
    int best_ch_layout = 0;
    int best_nb_channels   = 0;

    if (!codec->ch_layouts)
        return AV_CH_LAYOUT_STEREO;
    
    p = codec->ch_layouts;
    while (p->nb_channels != 0) {
        int nb_channels =  p->nb_channels;

        if (nb_channels > best_nb_channels) {
            best_ch_layout    = p->nb_channels;
            best_nb_channels = nb_channels;
        }
        p++;
    }
    return best_ch_layout;
}

int AudioEncoder::init(int bitRate, int channels, int sampleRate, int bitsPerSample, const char *aacFilePath, const char *codec_name) {
    
    _codecContext = NULL;
    _avFormatContext = NULL;
    _inputFrame = NULL;
    _samples = NULL;
    _sampleCursor = 0;
    
    _isWriteHeaderSuccess = false;
    
    _totalEncodeTimeMills = 0;
    _totalSWRTimeMills = 0;
    
    _publishBitRate = bitRate;
    _audioChannels = channels;
    _audioSampleRate = sampleRate;
    
    int ret;
    
    LOGI("aacFilePath is %s ", aacFilePath);
    avformat_network_init();
    
    // 1.协议
    const AVOutputFormat *oFormat = av_guess_format(NULL, aacFilePath, "audio");
    ret = avformat_alloc_output_context2(&_avFormatContext, oFormat, "aac", aacFilePath);
    if (ret < 0) {
        LOGI("avformat_alloc_output_context2 error");
        return ret;
    }

    // 打开文件
    ret = avio_open(&_avFormatContext->pb, aacFilePath, AVIO_FLAG_READ_WRITE);
    if (ret < 0) {
        LOGI("avio_open error");
        return ret;
    }
    
    // 2. 编码器
    // 2.1 查找编码器
    const AVCodec *c = avcodec_find_encoder_by_name("libfdk_aac");
    if (!c) {
        LOGI("avcodec_find_encoder_by_name error");
        return -1;
    }
    // 2.2 初始化编码器上下文
    _codecContext = avcodec_alloc_context3(c);
    
    // 2.3 设置编码器参数
    // 码率/比特率
    _codecContext->bit_rate = bitRate;
    // 编码格式
    _codecContext->sample_fmt = AV_SAMPLE_FMT_S16;
    if (!checkSampleFMT(c, AV_SAMPLE_FMT_S16)) {
        LOGI("no support sample fmt");
        return -1;
    }
    // 采样率
    _codecContext->sample_rate = bestSampleRate(c);
    
    //通道数
    _codecContext->channels = channels;
    _codecContext->channel_layout = channels == 1 ? AV_CH_LAYOUT_MONO : AV_CH_LAYOUT_STEREO;
    _codecContext->ch_layout.nb_channels = channels;
    
    //时间基
    _codecContext->time_base.num = 1;
    _codecContext->time_base.den = _codecContext->sample_rate;
    
    // 2.4 打开编码器
    ret = avcodec_open2(_codecContext, c, NULL);
    if (ret < 0) {
        LOGI("avcodec_open2 error");
        return -1;
    }
    
    // 必须要在avformat_write_header之前填充
    _audioStream = avformat_new_stream(_avFormatContext, NULL);
    _audioStream->id = 1;
    _audioStream->codecpar->codec_type = _codecContext->codec_type;
    _audioStream->codecpar->codec_id = _codecContext->codec_id;
    _audioStream->codecpar->ch_layout = _codecContext->ch_layout;
    _audioStream->codecpar->sample_rate = _codecContext->sample_rate;
    _audioStream->codecpar->bit_rate = _codecContext->bit_rate;

    _audioStream->codecpar->format = AV_SAMPLE_FMT_S16;
    _audioStream->codecpar->channels = channels;
    _audioStream->codecpar->channel_layout = _codecContext->channel_layout;
    
    av_dump_format(_avFormatContext, 0, aacFilePath, 1);
    
    ret = avformat_write_header(_avFormatContext, NULL);
    if (ret != 0) {
        LOGI("Could not write header\n");
        return -1;
    }
    _isWriteHeaderSuccess = true;

    return alloc_avframe();
}

int AudioEncoder::init(int bitRate, int channels, int bitesPerSample, const char *aacFilePath, const char *codec_name) {
    return init(bitRate, channels, 44100, bitesPerSample, aacFilePath, codec_name);
}

// 注意这里的存储数据的逻辑 每次只读frame_size的数据，多余的数据下次再读取
void AudioEncoder::encode(uint8_t *buffer, int size) {
    
    int bufferCursor = 0;
    int bufferSize = size;
    
    while (bufferSize >= (_bufferSize - _sampleCursor)) {
        int cpySize = _bufferSize - _sampleCursor;
        memcpy(_samples + _sampleCursor, buffer + bufferCursor, cpySize);
        bufferCursor += cpySize;
        bufferSize -= cpySize;
        if (encodePacket() < 0) {
            _sampleCursor = 0;
            break;
        }
        _sampleCursor = 0;
    }
    
    // 多余的数据
    if (bufferSize > 0) {
        memcpy(_samples + _sampleCursor, buffer + bufferCursor, bufferSize);
        _sampleCursor += bufferSize;
    }
}

void AudioEncoder::destory() {
    LOGI("start destroy!!!");

    if (_samples != NULL) {
        av_freep(&_samples);
        _samples = NULL;
    }
    
    if (_inputFrame != NULL) {
        av_frame_free(&_inputFrame);
        _inputFrame = NULL;
    }
    
    if (_outputPkt != NULL) {
        av_packet_free(&_outputPkt);
        _outputPkt = NULL;
    }
    
    if (_isWriteHeaderSuccess) {
        _avFormatContext->duration = _duration * AV_TIME_BASE;
        LOGI("duration is %.3f", this->_duration);
        av_write_trailer(_avFormatContext);
    }
    
    if (_avFormatContext->pb != NULL) {
        avio_close(_avFormatContext->pb);
    }
    
    if (_codecContext != NULL) {
        avcodec_close(_codecContext);
        av_free(_codecContext);
        _codecContext = NULL;
    }
    
    if (_avFormatContext != NULL) {
        avformat_free_context(_avFormatContext);
    }
    
    LOGI("end destroy!!! totalEncodeTimeMills is %d totalSWRTimeMills is %d", _totalEncodeTimeMills, _totalSWRTimeMills);
}

int AudioEncoder::encodePacket(){
    if (!_outputPkt) {
        _outputPkt = av_packet_alloc();
        if (!_outputPkt) {
            return -1;
        }
    }
    
    int ret = avcodec_send_frame(_codecContext, _inputFrame);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
        return 0;
    } else if(ret < 0) {
        return -1;
    }
    
    ret = avcodec_receive_packet(_codecContext, _outputPkt);
    if (ret == 0) {
        av_interleaved_write_frame(_avFormatContext, _outputPkt);
        av_packet_unref(_outputPkt);
        return 0;
    } else if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
        av_packet_unref(_outputPkt);
        return 0;
    } else {
        av_packet_unref(_outputPkt);
        return -1;
    }
}

// 初始化_inputFrame, 并绑定数据缓存区
int AudioEncoder::alloc_avframe() {
    int ret = 0;
    AVSampleFormat preferedSampleFMT = AV_SAMPLE_FMT_S16;
    int preferedChannels = _audioChannels;
    int preferedSampleRate = _audioSampleRate;
    _inputFrame = av_frame_alloc();
    if (!_inputFrame) {
        LOGI("Could not allocate audio frame\n");
        return -1;
    }
    
    _inputFrame->nb_samples = _codecContext->frame_size;
    _inputFrame->format = preferedSampleFMT;
    _inputFrame->channel_layout = preferedChannels == 1 ? AV_CH_LAYOUT_MONO : AV_CH_LAYOUT_STEREO;
    _inputFrame->ch_layout.nb_channels = preferedChannels;
    _inputFrame->sample_rate = preferedSampleRate;
    _bufferSize = av_samples_get_buffer_size(NULL, av_get_channel_layout_nb_channels(_inputFrame->channel_layout), _inputFrame->nb_samples, preferedSampleFMT, 0);
    
    _samples = (uint8_t *)av_malloc(_bufferSize);
    _sampleCursor = 0;
    if (!_samples) {
        LOGI("Could not allocate %d bytes for samples buffer\n", _bufferSize);
        return -2;
    }
    
    //填充frame音频数据和linesize指针
    ret = avcodec_fill_audio_frame(_inputFrame, av_get_channel_layout_nb_channels(_inputFrame->channel_layout), preferedSampleFMT, _samples, _bufferSize, 0);
    
    if (ret < 0) {
        LOGI("Could not setup audio frame\n");
    }
    return ret;
}
